#---------------------------------------------#
#   运行前一定要修改classes
#   如果生成的2007_train.txt里面没有目标信息
#   那么就是因为classes没有设定正确
#---------------------------------------------#

import xml.etree.ElementTree as ET
from os import getcwd

sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
#-----------------------------------------------------#
#   这里设定的classes顺序要和model_data里的txt一样
#-----------------------------------------------------#
classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]


# 作用：生成一个列表，该列表的每列数据均为“图片绝对路径 + (物体真实框的xywh + 类别序号)*n ”
# 提取bbox信息的操作不是固定的，你也可以把bbox的信息单独保存在一个txt文件中
def convert_annotation(year, image_id, list_file):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id), encoding='utf-8')
    tree=ET.parse(in_file)  # 解析xml文件
    root = tree.getroot()   # 获取根节点

    # 以<object>为单位遍历xmlTree，同时获取该物体的各种信息
    for obj in root.iter('object'):
        # 记录物体易识别程度
        difficult = 0 
        if obj.find('difficult')!=None:
            difficult = obj.find('difficult').text
        # 记录物体类别
        cls = obj.find('name').text
        # 不在训练的类中 或 识别难度高 则换下一个
        if cls not in classes or int(difficult)==1:
            continue
        # 找出对应类的下标
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        # 记录bbox的信息
        b = (int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)), int(float(xmlbox.find('xmax').text)), int(float(xmlbox.find('ymax').text)))
        list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))

# 记录当前目录
wd = getcwd().replace('\\', '/') 
# print(wd)

# VOC中的 训练、测试和验证都需经行bbox解析
for year, image_set in sets:
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set), encoding='utf-8').read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
    for image_id in image_ids:
        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg'%(wd, year, image_id))
        convert_annotation(year, image_id, list_file)
        list_file.write('\n')
    list_file.close()
